from pinecone import Pinecone, ServerlessSpec
from sklearn.datasets import load_digits
import matplotlib.pyplot as plt
import numpy as np
from collections import Counter

# 初始化Pinecone客户端
def initialize_pinecone(api_key):
    pinecone = Pinecone(api_key=api_key)
    return pinecone

# 管理索引：检查索引是否存在，如果不存在则创建新索引
def manage_index(pinecone, index_name):
    existing_indexes = pinecone.list_indexes()
    if any(index['name'] == index_name for index in existing_indexes): #索引存在
        print(f"索引 '{index_name}' 已存在。")
    else: #索引不存在
        print(f"索引 '{index_name}' 不存在。")
        print(f"正在创建新索引 '{index_name}'...")
        pinecone.create_index(
            name=index_name,
            dimension=64,  # MNIST 每个图像展平后是一个 64 维向量
            metric="euclidean",  # 使用欧氏距离
            spec=ServerlessSpec(cloud="aws", region="us-east-1")  # 服务器规格
        )
        print(f"索引 '{index_name}' 创建成功。")
    return pinecone.Index(index_name)

# 加载数据集并准备向量数据
def load_and_prepare_data():
    digits = load_digits(n_class=10)
    X = digits.data
    y = digits.target
    vectors = []
    for i, (image, label) in enumerate(zip(X, y)):
        vector_id = str(i)
        vector_values = image.tolist()
        metadata = {"label": int(label)}
        vectors.append((vector_id, vector_values, metadata))
    return vectors

def upload_data(index, vectors, batch_size=1000):
    # 执行一个查询来检查索引中是否存在向量
    test_query_vector = [0] * 64  # 假设索引的维度是64，根据实际情况修改
    results = index.query(vector=test_query_vector, top_k=1, include_metadata=False)
    
    # 如果查询结果为空，说明索引中没有数据
    if not results['matches']:
        print("索引为空，开始上传数据...")
        for i in range(0, len(vectors), batch_size):
            batch = vectors[i:i + batch_size]
            index.upsert(batch)
    else:
        print("索引中已存在数据，跳过上传。")
# 创建查询向量
def create_query_vector():
    digit_3 = np.array([
        [0, 0, 255, 255, 255, 255, 0, 0],
        [0, 0, 0, 0, 0, 255, 0, 0],
        [0, 0, 0, 0, 0, 255, 0, 0],
        [0, 0, 0, 255, 255, 255, 0, 0],
        [0, 0, 0, 0, 0, 255, 0, 0],
        [0, 0, 0, 0, 0, 255, 0, 0],
        [0, 0, 0, 0, 0, 255, 0, 0],
        [0, 0, 255, 255, 255, 255, 0, 0]
    ])
    digit_3_flatten = (digit_3 / 255.0) * 16
    query_data = digit_3_flatten.ravel().tolist()
    return digit_3, query_data

# 在Pinecone索引中执行查询并预测结果
def query_and_predict(index, query_data, top_k=11):
    results = index.query(vector=query_data, top_k=top_k, include_metadata=True)
    labels = [match['metadata']['label'] for match in results['matches']]
    # 打印每个匹配结果的详细信息
    for match, label in zip(results['matches'], labels):
        print(f"id: {match['id']}, distance: {match['score']}, label: {label}")
    if labels:
        final_prediction = Counter(labels).most_common(1)[0][0]
    else:
        final_prediction = None
    return final_prediction, results

# 显示查询图像和预测结果
def display_result(digit_3, final_prediction):
    if final_prediction is not None:
        plt.imshow(digit_3, cmap='gray')
        plt.title(f"Predicted digit: {final_prediction}", size=15)
        plt.axis('off')
        plt.show()
    else:
        print("没有找到匹配的结果。")

# 主函数，程序入口点
def main(api_key, index_name):
    pinecone = initialize_pinecone(api_key)
    index = manage_index(pinecone, index_name)
    vectors = load_and_prepare_data()
    upload_data(index, vectors)
    digit_3, query_data = create_query_vector()
    final_prediction, results = query_and_predict(index, query_data)
    display_result(digit_3, final_prediction)

# 检查脚本是否作为主程序运行
if __name__ == "__main__":
    
    main("0f4bb63e-9728-43d4-89b8-562f96acb31b", "mnist-index")
    